import os
from PIL import Image
from transformers import CLIPProcessor, CLIPModel
import shutil

# 初始化CLIP模型和处理器
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

def classify_image(image_path, candidate_labels):
    image = Image.open(image_path).convert("RGB")
    inputs = processor(text=candidate_labels, images=image, return_tensors="pt", padding=True)
    outputs = model(**inputs)
    logits_per_image = outputs.logits_per_image  # 图片对标签的相似度
    probs = logits_per_image.softmax(dim=1)
    max_prob, idx = probs.max(dim=1)
    return candidate_labels[idx], max_prob.item()

def select_top_images(folder_path, candidate_labels, top_k=9):
    scored_images = []
    for filename in os.listdir(folder_path):
        if filename.lower().endswith((".jpg", ".png", ".jpeg")):
            path = os.path.join(folder_path, filename)
            label, score = classify_image(path, candidate_labels)
            scored_images.append((path, label, score))
    # 按分数排序
    scored_images.sort(key=lambda x: x[2], reverse=True)
    # 选出top_k张
    return scored_images[:top_k]

if __name__ == "__main__":
    folder = "../sample-pics"
    labels = ["person", "scenery", "animal"]
    top_images = select_top_images(folder, labels, top_k=3)
    output_dir = "./selected_images"
    os.makedirs(output_dir, exist_ok=True)
    for img_path, label, score in top_images:
        print(f"Selected {os.path.basename(img_path)} as {label} with score {score:.4f}")
        shutil.copy(img_path, output_dir)